package ezai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"gitee.com/hongzhaomin/ezai/bean"
	"gitee.com/hongzhaomin/gspring/iface"
	"gitee.com/hongzhaomin/hzm-common-go/beautify"
	"gitee.com/hongzhaomin/hzm-common-go/toolkit"
	"github.com/tmaxmax/go-sse"
	"io"
	"net/http"
	"strings"
)

var _ ChatClient = (*DefaultChatClient)(nil)

type DefaultChatClient struct {
	ConfigBean  *ConfigBean         `@Resource:""`
	Logger      iface.GspringLogger `@Resource:""`
	ChatSession ChatSession         `@Resource:""`
	ToolParser  *ToolParser         `@Resource:""`
}

func (my *DefaultChatClient) ChatStream(sessionId string, userContent string) <-chan string {
	ch := make(chan string, 5)
	message := bean.Message{
		Role:    RoleUser,
		Content: userContent,
	}
	go func() {
		my.chat(sessionId, message, func(output string) {
			ch <- output
		})
		// 关闭channel
		close(ch)
	}()
	return ch
}

func (my *DefaultChatClient) Chat(sessionId, userContent string) string {
	output := ""
	message := bean.Message{
		Role:    RoleUser,
		Content: userContent,
	}
	my.chat(sessionId, message, func(output string) {})
	return output
}

func (my *DefaultChatClient) chat(sessionId string, message bean.Message, fn func(output string)) {
	msgObjs := my.ChatSession.GetMessages(sessionId)
	if msgObjs == nil {
		sysMsg := sysMsgMap[toolSystemContent]
		my.ChatSession.AddMessage(sessionId, &sysMsg)
	}
	my.ChatSession.AddMessage(sessionId, &message)

	assistantMsg := &bytes.Buffer{}
	var toolCall *bean.ToolCall
	var deltaTool *bean.Message
	my.ChatCreate(sessionId, func(res string) {
		var aiResp bean.AiResp
		err := json.Unmarshal([]byte(res), &aiResp)
		if err != nil {
			my.Logger.Error("Unmarshal aiResp fail: response: %s, %v", res, err)
			return
		}
		choice := aiResp.Choices[0]

		delta := choice.Delta
		if !my.ConfigBean.Stream {
			delta = choice.Message
		}

		// 智谱大模型（GLM-Z1-Flash）可能function_call有点问题，它把tool_calls放到content里面了，这里处理一下
		if len(delta.ToolCalls) == 0 && strings.Contains(delta.Content, "tool_calls") {
			split := strings.Split(delta.Content, "\n")
			lastIndex := len(split) - 1
			for i, s := range split {
				if i == lastIndex {
					choiceJson := s
					var ch bean.Choice
					err = json.Unmarshal([]byte(choiceJson), &ch)
					if err != nil {
						my.Logger.Error("Unmarshal choice fail: response: %s, %v", res, err)
					} else {
						choice = ch
						delta = choice.Delta
					}
				}
			}
		}

		if choice.FinishReason == "tool_calls" {
			if len(delta.ToolCalls) > 0 {
				toolCall = &delta.ToolCalls[0]
				deltaTool = &delta
				return
			}
		}
		assistantMsg.WriteString(delta.Content)
		fn(delta.Content)
	})

	result := assistantMsg.String()
	if result != "" {
		my.ChatSession.AddMessage(sessionId, &bean.Message{
			Role:    RoleAssistant,
			Content: result,
		})
	}

	if toolCall == nil {
		return
	}
	fn(fmt.Sprintf("\n正在调用函数：%s\n%s\n", toolCall.Function.Name, beautify.FmtJson(toolCall)))
	// 执行函数调用
	callToolMsg := my.callTool(*toolCall)
	if deltaTool != nil {
		// 把模型输出放入消息
		my.ChatSession.AddMessage(sessionId, deltaTool)
	}
	my.chat(sessionId, callToolMsg, fn)
	return
}

func (my *DefaultChatClient) ChatCreate(sessionId string, fn func(res string)) {
	req, cancelFunc := my.buildReq(sessionId)
	if !my.ConfigBean.Stream {
		fn(my.httpRequest(req))
		cancelFunc()
	} else {
		my.sseRequest(req, cancelFunc, fn)
	}
}

func (my *DefaultChatClient) httpRequest(req *http.Request) string {
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		my.Logger.Error("模型访问异常：%v", err)
		return ""
	}
	all, err := io.ReadAll(res.Body)
	if res.StatusCode != http.StatusOK {
		my.Logger.Error("模型访问失败：%s", string(all))
		return ""
	}
	if err != nil {
		my.Logger.Error("读取模型返回数据异常：%v", err)
		return ""
	}
	return string(all)
}

func (my *DefaultChatClient) sseRequest(req *http.Request, cancelFunc context.CancelFunc, fn func(res string)) {
	// 定义事件处理函数
	var remover sse.EventCallbackRemover
	handler := func(event sse.Event) {
		if event.Data != "" {
			//my.Logger.Info("收到数据: %s\n", event.Data)
			if event.Data == "[DONE]" {
				cancelFunc()
				return
			}
			fn(event.Data)
		}
	}

	// 订阅事件流（支持自动重连）
	conn := sse.NewConnection(req)
	remover = conn.SubscribeMessages(handler)
	if err := conn.Connect(); err != nil {
		if err.Error() != "context canceled" {
			my.Logger.Error("连接异常：%v", err)
		}
	}
	remover()
}

func (my *DefaultChatClient) buildReq(sessionId string) (*http.Request, context.CancelFunc) {
	params := bean.AiReq{
		Model:    my.ConfigBean.Model,
		Messages: my.ChatSession.GetMessages(sessionId),
		Stream:   my.ConfigBean.Stream,
	}
	tools := my.ToolParser.tools
	if len(tools) > 0 {
		var v = false
		params.DoSample = &v
		params.Tools = tools
	}
	paramsJson := toolkit.ToJson3(params)
	my.Logger.Info("请求模型参数：%s", paramsJson)
	body := bytes.NewBufferString(paramsJson)

	// 创建 HTTP 请求
	ctx, cancelFunc := context.WithCancel(context.Background())
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, my.ConfigBean.Url, body)
	if err != nil {
		my.Logger.Error("构建请求参数异常：%v", err)
		cancelFunc()
		return nil, nil
	}

	// 设置请求头
	req.Header.Set("Content-Type", "application/json;charset=utf-8")
	req.Header.Set("Authorization", my.ConfigBean.ApiKey)
	return req, cancelFunc
}

func (my *DefaultChatClient) callTool(tool bean.ToolCall) bean.Message {
	resultMsg := bean.Message{
		Role:       RoleTool,
		Content:    "未查询到任何有用信息",
		ToolCallId: tool.Id,
	}
	function := tool.Function
	val, err := my.ToolParser.FunctionCall(*function)
	if err != nil {
		my.Logger.Error("工具[%s]调用失败：%v", function.Name, err)
		return resultMsg
	}
	if s, ok := val.(string); ok {
		resultMsg.Content = s
	} else {
		resultMsg.Content, err = toolkit.ConvertStr4Any(val)
		if err != nil {
			resultMsg.Content = toolkit.ToJson3(val)
		}
	}
	return resultMsg
}
